package org.jeecg.common.util;

import org.bouncycastle.asn1.gm.GMNamedCurves;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.crypto.engines.SM2Engine;
import org.bouncycastle.crypto.generators.ECKeyPairGenerator;
import org.bouncycastle.crypto.params.*;
import org.bouncycastle.crypto.signers.SM2Signer;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPrivateKey;
import org.bouncycastle.jcajce.provider.asymmetric.ec.BCECPublicKey;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.jce.spec.ECNamedCurveParameterSpec;
import org.bouncycastle.jce.spec.ECParameterSpec;
import org.bouncycastle.util.encoders.Base64;

import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * SM2 加密工具类
 * 提供 SM2 密钥生成、加密、解密、签名、验签功能
 * 依赖 Bouncy Castle 库
 */
public class SM2Utils {
    static {
        // 注册 Bouncy Castle 安全提供者
        if (Security.getProvider(BouncyCastleProvider.PROVIDER_NAME) == null) {
            Security.addProvider(new BouncyCastleProvider());
        }
    }

    // SM2 曲线名称
    private static final String SM2_CURVE_NAME = "sm2p256v1";

    // 编码格式
    private static final String CHARSET = "UTF-8";

    /**
     * 生成 SM2 密钥对
     * @return 密钥对对象
     * @throws Exception 生成失败时抛出异常
     */
    public static KeyPair generateKeyPair() throws Exception {
        X9ECParameters sm2Parameters = GMNamedCurves.getByName(SM2_CURVE_NAME);
        ECDomainParameters domainParameters = new ECDomainParameters(
                sm2Parameters.getCurve(),
                sm2Parameters.getG(),
                sm2Parameters.getN()
        );

        ECKeyPairGenerator keyPairGenerator = new ECKeyPairGenerator();
        keyPairGenerator.init(new ECKeyGenerationParameters(domainParameters, new SecureRandom()));
        AsymmetricCipherKeyPair keyPair = keyPairGenerator.generateKeyPair();

        ECPrivateKeyParameters privateKeyParams = (ECPrivateKeyParameters) keyPair.getPrivate();
        ECPublicKeyParameters publicKeyParams = (ECPublicKeyParameters) keyPair.getPublic();

        ECParameterSpec parameterSpec = new ECNamedCurveParameterSpec(
                SM2_CURVE_NAME,
                sm2Parameters.getCurve(),
                sm2Parameters.getG(),
                sm2Parameters.getN()
        );
        PublicKey publicKey = new BCECPublicKey("EC", publicKeyParams, parameterSpec, BouncyCastleProvider.CONFIGURATION);

        PrivateKey privateKey = new BCECPrivateKey("EC", privateKeyParams, (BCECPublicKey) publicKey, parameterSpec, BouncyCastleProvider.CONFIGURATION);
        //PublicKey publicKey = new BCECPublicKey("EC", publicKeyParams, parameterSpec, BouncyCastleProvider.CONFIGURATION);

        return new KeyPair(publicKey, privateKey);
    }

    /**
     * SM2 加密（C1C3C2 格式）
     * @param plainText 明文
     * @param publicKey 公钥
     * @return Base64 编码的密文
     * @throws Exception 加密失败时抛出异常
     */
    public static String encrypt(String plainText, PublicKey publicKey) throws Exception {
        BCECPublicKey ecPublicKey = (BCECPublicKey) publicKey;
        ECPublicKeyParameters ecPublicKeyParameters = new ECPublicKeyParameters(
                ecPublicKey.getQ(),
                new ECDomainParameters(
                        ecPublicKey.getParameters().getCurve(),
                        ecPublicKey.getParameters().getG(),
                        ecPublicKey.getParameters().getN()
                )
        );

        SM2Engine sm2Engine = new SM2Engine();
        sm2Engine.init(true, new ParametersWithRandom(ecPublicKeyParameters, new SecureRandom()));

        byte[] encryptedBytes = sm2Engine.processBlock(plainText.getBytes(CHARSET), 0, plainText.getBytes(CHARSET).length);
        return Base64.toBase64String(encryptedBytes);
    }

    /**
     * SM2 解密（C1C3C2 格式）
     * @param cipherText Base64 编码的密文
     * @param privateKey 私钥
     * @return 明文
     * @throws Exception 解密失败时抛出异常
     */
    public static String decrypt(String cipherText, PrivateKey privateKey) throws Exception {
        BCECPrivateKey ecPrivateKey = (BCECPrivateKey) privateKey;
        ECPrivateKeyParameters ecPrivateKeyParameters = new ECPrivateKeyParameters(
                ecPrivateKey.getD(),
                new ECDomainParameters(
                        ecPrivateKey.getParameters().getCurve(),
                        ecPrivateKey.getParameters().getG(),
                        ecPrivateKey.getParameters().getN()
                )
        );

        SM2Engine sm2Engine = new SM2Engine();
        sm2Engine.init(false, ecPrivateKeyParameters);

        byte[] decryptedBytes = sm2Engine.processBlock(Base64.decode(cipherText), 0, Base64.decode(cipherText).length);
        return new String(decryptedBytes, CHARSET);
    }

    /**
     * SM2 签名
     * @param message 待签名消息
     * @param privateKey 私钥
     * @return Base64 编码的签名
     * @throws Exception 签名失败时抛出异常
     */
    public static String sign(String message, PrivateKey privateKey) throws Exception {
        BCECPrivateKey ecPrivateKey = (BCECPrivateKey) privateKey;
        ECPrivateKeyParameters privateKeyParameters = new ECPrivateKeyParameters(
                ecPrivateKey.getD(),
                new ECDomainParameters(
                        ecPrivateKey.getParameters().getCurve(),
                        ecPrivateKey.getParameters().getG(),
                        ecPrivateKey.getParameters().getN()
                )
        );

        SM2Signer signer = new SM2Signer();
        signer.init(true, new ParametersWithRandom(privateKeyParameters, new SecureRandom()));
        signer.update(message.getBytes(CHARSET), 0, message.getBytes(CHARSET).length);

        byte[] signature = signer.generateSignature();
        return Base64.toBase64String(signature);
    }

    /**
     * SM2 验签
     * @param message 原始消息
     * @param signatureBase64 Base64 编码的签名
     * @param publicKey 公钥
     * @return 验签结果
     * @throws Exception 验签失败时抛出异常
     */
    public static boolean verify(String message, String signatureBase64, PublicKey publicKey) throws Exception {
        BCECPublicKey ecPublicKey = (BCECPublicKey) publicKey;
        ECPublicKeyParameters publicKeyParameters = new ECPublicKeyParameters(
                ecPublicKey.getQ(),
                new ECDomainParameters(
                        ecPublicKey.getParameters().getCurve(),
                        ecPublicKey.getParameters().getG(),
                        ecPublicKey.getParameters().getN()
                )
        );

        SM2Signer signer = new SM2Signer();
        signer.init(false, publicKeyParameters);
        signer.update(message.getBytes(CHARSET), 0, message.getBytes(CHARSET).length);

        return signer.verifySignature(Base64.decode(signatureBase64));
    }

    /**
     * 公钥转 Base64
     * @param publicKey 公钥
     * @return Base64 编码的公钥
     */
    public static String publicKeyToBase64(PublicKey publicKey) {
        return Base64.toBase64String(publicKey.getEncoded());
    }

    /**
     * 私钥转 Base64
     * @param privateKey 私钥
     * @return Base64 编码的私钥
     */
    public static String privateKeyToBase64(PrivateKey privateKey) {
        return Base64.toBase64String(privateKey.getEncoded());
    }

    /**
     * Base64 转公钥
     * @param publicKeyBase64 Base64 编码的公钥
     * @return 公钥对象
     * @throws Exception 转换失败时抛出异常
     */
    public static PublicKey base64ToPublicKey(String publicKeyBase64) throws Exception {
        byte[] keyBytes = Base64.decode(publicKeyBase64);
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("EC", BouncyCastleProvider.PROVIDER_NAME);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * Base64 转私钥
     * @param privateKeyBase64 Base64 编码的私钥
     * @return 私钥对象
     * @throws Exception 转换失败时抛出异常
     */
    public static PrivateKey base64ToPrivateKey(String privateKeyBase64) throws Exception {
        byte[] keyBytes = Base64.decode(privateKeyBase64);
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
        KeyFactory keyFactory = KeyFactory.getInstance("EC", BouncyCastleProvider.PROVIDER_NAME);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * 从 DER 格式公钥字节转换为 PublicKey
     * @param derPublicKey DER 格式公钥字节
     * @return 公钥对象
     * @throws Exception 转换失败时抛出异常
     */
    public static PublicKey derToPublicKey(byte[] derPublicKey) throws Exception {
        X509EncodedKeySpec keySpec = new X509EncodedKeySpec(derPublicKey);
        KeyFactory keyFactory = KeyFactory.getInstance("EC", BouncyCastleProvider.PROVIDER_NAME);
        return keyFactory.generatePublic(keySpec);
    }

    /**
     * 从 DER 格式私钥字节转换为 PrivateKey
     * @param derPrivateKey DER 格式私钥字节
     * @return 私钥对象
     * @throws Exception 转换失败时抛出异常
     */
    public static PrivateKey derToPrivateKey(byte[] derPrivateKey) throws Exception {
        PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(derPrivateKey);
        KeyFactory keyFactory = KeyFactory.getInstance("EC", BouncyCastleProvider.PROVIDER_NAME);
        return keyFactory.generatePrivate(keySpec);
    }

    /**
     * 提取公钥的原始字节（不包含 DER 头部）
     * @param publicKey 公钥对象
     * @return 原始公钥字节
     */
    public static byte[] getRawPublicKey(PublicKey publicKey) {
        BCECPublicKey ecPublicKey = (BCECPublicKey) publicKey;
        return ecPublicKey.getQ().getEncoded(false);
    }

    /**
     * 从原始字节构建公钥
     * @param rawPublicKey 原始公钥字节
     * @return 公钥对象
     * @throws Exception 构建失败时抛出异常
     */
    public static PublicKey buildPublicKeyFromRaw(byte[] rawPublicKey) throws Exception {
        X9ECParameters sm2Parameters = GMNamedCurves.getByName(SM2_CURVE_NAME);
        ECDomainParameters domainParameters = new ECDomainParameters(
                sm2Parameters.getCurve(),
                sm2Parameters.getG(),
                sm2Parameters.getN()
        );

        org.bouncycastle.math.ec.ECPoint point = domainParameters.getCurve().decodePoint(rawPublicKey);
        ECPublicKeyParameters publicKeyParameters = new ECPublicKeyParameters(point, domainParameters);

        ECParameterSpec parameterSpec = new ECNamedCurveParameterSpec(
                SM2_CURVE_NAME,
                sm2Parameters.getCurve(),
                sm2Parameters.getG(),
                sm2Parameters.getN()
        );

        return new BCECPublicKey("EC", publicKeyParameters, parameterSpec, BouncyCastleProvider.CONFIGURATION);
    }
    public static void main(String[] args) {
        try {
            // 1. 生成密钥对
            KeyPair keyPair = SM2Utils.generateKeyPair();
            PublicKey publicKey = keyPair.getPublic();
            PrivateKey privateKey = keyPair.getPrivate();

            // 2. 转换为 Base64 存储或传输
            String publicKeyBase64 = SM2Utils.publicKeyToBase64(publicKey);
            String privateKeyBase64 = SM2Utils.privateKeyToBase64(privateKey);
            System.out.println("公钥(Base64): " + publicKeyBase64);
            System.out.println("私钥(Base64): " + privateKeyBase64);

            // 3. 从 Base64 恢复密钥
            PublicKey restoredPublicKey = SM2Utils.base64ToPublicKey(publicKeyBase64);
            PrivateKey restoredPrivateKey = SM2Utils.base64ToPrivateKey(privateKeyBase64);

            // 4. 加密和解密示例
            String plainText = "Hello, SM2 加密示例!";
            String encrypted = SM2Utils.encrypt(plainText, restoredPublicKey);
            System.out.println("加密后: " + encrypted);

            String decrypted = SM2Utils.decrypt(encrypted, restoredPrivateKey);
            System.out.println("解密后: " + decrypted);

            // 5. 签名和验签示例
            String message = "需要签名的消息内容";
            String signature = SM2Utils.sign(message, restoredPrivateKey);
            System.out.println("签名: " + signature);

            boolean verifyResult = SM2Utils.verify(message, signature, restoredPublicKey);
            System.out.println("验签结果: " + verifyResult);

            // 6. 原始公钥处理示例
            byte[] rawPublicKey = SM2Utils.getRawPublicKey(restoredPublicKey);
            PublicKey rebuiltPublicKey = SM2Utils.buildPublicKeyFromRaw(rawPublicKey);
            System.out.println("重建公钥验证: " + rebuiltPublicKey.equals(restoredPublicKey));

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}